#!/usr/bin/env python3
"""
多线程压力测试 - 分层次测试并发能力
测试你的数据库系统在不同并发压力下的表现
"""

import os
import sys
import time
import tempfile
import shutil
import random
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import List, Dict, Any, Tuple
import queue
import psutil  # 需要安装: pip install psutil

# 路径设置
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)

try:
    from database_server_adapter import DatabaseServerAdapter
    from table_schema import Column, ColumnType
    from storage_engine import CachePolicy
except ImportError as e:
    print(f"导入失败: {e}")
    sys.exit(1)


@dataclass
class StressTestResult:
    """压力测试结果"""
    level: str
    thread_count: int
    operations_per_thread: int
    total_operations: int
    success_count: int
    failure_count: int
    total_time: float
    avg_response_time: float
    operations_per_second: float
    memory_usage_mb: float
    cpu_usage_percent: float
    errors: List[str]


class MultithreadingStressTest:
    """多线程压力测试类"""
    
    def __init__(self):
        self.temp_dir = None
        self.adapter = None
        self.results = []
        self.test_data_generator = self._create_test_data_generator()
        
    def setup(self):
        """初始化测试环境"""
        self.temp_dir = tempfile.mkdtemp(prefix="stress_test_")
        self.adapter = DatabaseServerAdapter(
            data_root=self.temp_dir,
            cache_size=200,  # 增大缓存以应对高并发
            cache_policy=CachePolicy.LRU,
            enable_locking=True,
            max_query_workers=8,  # 增加工作线程
            max_io_workers=4
        )
        
        # 认证
        if not self.adapter.login("root", "root12345"):
            raise Exception("认证失败")
        
        # 创建测试数据库和表
        self.adapter.create_database("stress_test_db")
        self.adapter.use_database("stress_test_db")
        
        # 创建用户表
        user_columns = [
            Column("user_id", ColumnType.INT, primary_key=True),
            Column("username", ColumnType.VARCHAR, length=50),
            Column("email", ColumnType.VARCHAR, length=100),
            Column("age", ColumnType.INT),
            Column("balance", ColumnType.FLOAT),
            Column("is_active", ColumnType.BOOLEAN),
            Column("created_at", ColumnType.VARCHAR, length=20)
        ]
        self.adapter.create_table("users", user_columns)
        
        # 创建订单表
        order_columns = [
            Column("order_id", ColumnType.INT, primary_key=True),
            Column("user_id", ColumnType.INT),
            Column("product_name", ColumnType.VARCHAR, length=100),
            Column("amount", ColumnType.FLOAT),
            Column("quantity", ColumnType.INT),
            Column("status", ColumnType.VARCHAR, length=20)
        ]
        self.adapter.create_table("orders", order_columns)
        
        print("✅ 测试环境初始化完成")
    
    def cleanup(self):
        """清理测试环境"""
        if self.adapter:
            self.adapter.logout()
            self.adapter.close()
        
        if self.temp_dir and os.path.exists(self.temp_dir):
            try:
                shutil.rmtree(self.temp_dir)
            except Exception as e:
                print(f"清理失败: {e}")
    
    def _create_test_data_generator(self):
        """创建测试数据生成器"""
        def generate_user(user_id):
            return {
                "user_id": user_id,
                "username": f"user_{user_id}",
                "email": f"user{user_id}@test.com",
                "age": random.randint(18, 65),
                "balance": round(random.uniform(0, 10000), 2),
                "is_active": random.choice([True, False]),
                "created_at": f"2024-{random.randint(1,12):02d}-{random.randint(1,28):02d}"
            }
        
        def generate_order(order_id, user_id):
            products = ["笔记本电脑", "手机", "平板", "键盘", "鼠标", "耳机", "音响", "摄像头"]
            return {
                "order_id": order_id,
                "user_id": user_id,
                "product_name": random.choice(products),
                "amount": round(random.uniform(50, 5000), 2),
                "quantity": random.randint(1, 10),
                "status": random.choice(["pending", "paid", "shipped", "delivered"])
            }
        
        return {"user": generate_user, "order": generate_order}
    
    def _get_system_metrics(self):
        """获取系统指标"""
        process = psutil.Process()
        memory_mb = process.memory_info().rss / 1024 / 1024
        cpu_percent = process.cpu_percent()
        return memory_mb, cpu_percent
    
    def _worker_mixed_operations(self, worker_id: int, operations_count: int, 
                                result_queue: queue.Queue):
        """工作线程 - 执行混合操作"""
        success_count = 0
        failure_count = 0
        errors = []
        response_times = []
        
        # 使用共享的适配器，避免多个线程同时初始化系统数据库
        worker_adapter = self.adapter
        
        try:
            # 不需要重新认证，使用主线程已认证的连接
            # 确保使用正确的数据库
            if not worker_adapter.get_current_database():
                worker_adapter.use_database("stress_test_db")
            
            for i in range(operations_count):
                start_time = time.time()
                
                try:
                    # 随机选择操作类型
                    operation_type = random.choices(
                        ["insert", "select", "update", "delete"],
                        weights=[40, 35, 15, 10]  # 插入和查询权重较高
                    )[0]
                    
                    if operation_type == "insert":
                        # 插入操作
                        base_id = worker_id * 10000 + i
                        if random.choice([True, False]):
                            # 插入用户
                            user_data = self.test_data_generator["user"](base_id)
                            result = worker_adapter.insert_record("users", user_data)
                        else:
                            # 插入订单
                            user_id = random.randint(1, 1000)  # 引用可能存在的用户
                            order_data = self.test_data_generator["order"](base_id, user_id)
                            result = worker_adapter.insert_record("orders", order_data)
                    
                    elif operation_type == "select":
                        # 查询操作
                        table = random.choice(["users", "orders"])
                        if random.choice([True, False]):
                            # 条件查询
                            if table == "users":
                                conditions = [
                                    {"age": {"$gt": 30}},
                                    {"is_active": True},
                                    {"balance": {"$gt": 1000}},
                                    {"user_id": {"$lte": 100}}
                                ]
                            else:
                                conditions = [
                                    {"status": "paid"},
                                    {"amount": {"$gt": 500}},
                                    {"quantity": {"$gte": 2}},
                                    {"user_id": {"$lte": 100}}
                                ]
                            condition = random.choice(conditions)
                            result = worker_adapter.select_records(table, condition)
                        else:
                            # 全表查询
                            result = worker_adapter.select_records(table)
                        result = result is not None
                    
                    elif operation_type == "update":
                        # 更新操作
                        table = random.choice(["users", "orders"])
                        if table == "users":
                            condition = {"age": {"$lt": 25}}
                            updates = {"is_active": True}
                        else:
                            condition = {"status": "pending"}
                            updates = {"status": "confirmed"}
                        result = worker_adapter.update_records(table, condition, updates)
                        result = result >= 0  # 更新成功（包括0条）
                    
                    elif operation_type == "delete":
                        # 删除操作（谨慎删除）
                        table = random.choice(["users", "orders"])
                        if table == "users":
                            condition = {"balance": {"$lt": 10}}
                        else:
                            condition = {"status": "cancelled"}
                        result = worker_adapter.delete_records(table, condition)
                        result = result >= 0  # 删除成功（包括0条）
                    
                    response_time = time.time() - start_time
                    response_times.append(response_time)
                    
                    if result:
                        success_count += 1
                    else:
                        failure_count += 1
                        errors.append(f"Worker {worker_id} 操作 {i+1} 失败: {operation_type}")
                
                except Exception as e:
                    failure_count += 1
                    response_time = time.time() - start_time
                    response_times.append(response_time)
                    errors.append(f"Worker {worker_id} 操作 {i+1} 异常: {str(e)}")
                
                # 随机延迟（模拟真实场景）
                if random.random() < 0.1:  # 10%概率短暂延迟
                    time.sleep(random.uniform(0.001, 0.01))
        
        finally:
            # 不需要关闭共享连接
            pass
        
        result_queue.put((worker_id, success_count, failure_count, response_times, errors))
    
    def run_stress_test_level(self, level_name: str, thread_count: int, 
                            operations_per_thread: int) -> StressTestResult:
        """运行指定级别的压力测试"""
        print(f"\n🔥 {level_name} - {thread_count}线程，每线程{operations_per_thread}操作")
        print("-" * 50)
        
        # 记录初始系统状态
        initial_memory, initial_cpu = self._get_system_metrics()
        
        result_queue = queue.Queue()
        start_time = time.time()
        
        # 启动工作线程
        threads = []
        for i in range(thread_count):
            thread = threading.Thread(
                target=self._worker_mixed_operations,
                args=(i, operations_per_thread, result_queue)
            )
            threads.append(thread)
            thread.start()
        
        # 等待所有线程完成
        for thread in threads:
            thread.join()
        
        total_time = time.time() - start_time
        
        # 收集结果
        total_success = 0
        total_failure = 0
        all_response_times = []
        all_errors = []
        
        while not result_queue.empty():
            worker_id, success, failure, response_times, errors = result_queue.get()
            total_success += success
            total_failure += failure
            all_response_times.extend(response_times)
            all_errors.extend(errors)
        
        # 计算指标
        total_operations = thread_count * operations_per_thread
        avg_response_time = sum(all_response_times) / len(all_response_times) if all_response_times else 0
        operations_per_second = total_operations / total_time if total_time > 0 else 0
        
        # 获取峰值系统指标
        final_memory, final_cpu = self._get_system_metrics()
        
        result = StressTestResult(
            level=level_name,
            thread_count=thread_count,
            operations_per_thread=operations_per_thread,
            total_operations=total_operations,
            success_count=total_success,
            failure_count=total_failure,
            total_time=total_time,
            avg_response_time=avg_response_time,
            operations_per_second=operations_per_second,
            memory_usage_mb=max(initial_memory, final_memory),
            cpu_usage_percent=max(initial_cpu, final_cpu),
            errors=all_errors[:10]  # 只保留前10个错误
        )
        
        # 打印结果
        print(f"   📊 操作统计: {total_success}/{total_operations} 成功 ({total_success/total_operations*100:.1f}%)")
        print(f"   ⏱️  总耗时: {total_time:.2f}秒")
        print(f"   🚀 吞吐量: {operations_per_second:.1f} ops/sec")
        print(f"   ⚡ 平均响应时间: {avg_response_time*1000:.2f}ms")
        print(f"   💾 内存使用: {final_memory:.1f}MB")
        print(f"   🔥 CPU使用: {final_cpu:.1f}%")
        
        if all_errors:
            print(f"   ⚠️  错误数: {len(all_errors)} (显示前3个)")
            for error in all_errors[:3]:
                print(f"      {error}")
        
        return result
    
    def run_all_stress_levels(self):
        """运行所有压力级别的测试"""
        print("🚀 多线程压力测试 - 分层次压力测试")
        print("=" * 60)
        
        # 定义压力级别
        stress_levels = [
            ("轻压力", 5, 50),      # 5线程 × 50操作 = 250操作
            ("中压力", 10, 100),    # 10线程 × 100操作 = 1000操作
            ("高压力", 20, 100),    # 20线程 × 100操作 = 2000操作
            ("极限压力", 50, 100),   # 50线程 × 100操作 = 5000操作
            ("疯狂压力", 100, 50),   # 100线程 × 50操作 = 5000操作
            ("超级压力", 200, 25),   # 200线程 × 25操作 = 5000操作
        ]
        
        for level_name, thread_count, ops_per_thread in stress_levels:
            try:
                result = self.run_stress_test_level(level_name, thread_count, ops_per_thread)
                self.results.append(result)
                
                # 如果成功率低于80%，停止更高级别的测试
                success_rate = result.success_count / result.total_operations
                if success_rate < 0.8:
                    print(f"\n⚠️  {level_name}成功率{success_rate*100:.1f}%过低，停止更高级别测试")
                    break
                
                # 让系统稍作休息
                time.sleep(2)
                
            except Exception as e:
                print(f"\n❌ {level_name}测试异常: {e}")
                break
    
    def generate_final_report(self):
        """生成最终测试报告"""
        print("\n" + "=" * 60)
        print("📊 多线程压力测试最终报告")
        print("=" * 60)
        
        if not self.results:
            print("❌ 没有测试结果")
            return
        
        # 找出系统极限
        if self.results:
            max_ops_per_sec = max(result.operations_per_second for result in self.results)
            valid_results = [result for result in self.results 
                           if result.success_count / result.total_operations >= 0.9]
            max_concurrent_threads = max(result.thread_count for result in valid_results) if valid_results else 0
        else:
            max_ops_per_sec = 0
            max_concurrent_threads = 0
        
        print("🏆 系统性能极限:")
        print(f"   📈 最大吞吐量: {max_ops_per_sec:.1f} operations/second")
        print(f"   🧵 最大并发线程: {max_concurrent_threads} threads (成功率≥90%)")
        
        # 详细结果表
        print("\n📋 详细测试结果:")
        print(f"{'级别':<12} {'线程数':<8} {'成功率':<8} {'吞吐量':<12} {'响应时间':<12} {'内存':<10}")
        print("-" * 70)
        
        for result in self.results:
            success_rate = result.success_count / result.total_operations * 100
            print(f"{result.level:<12} {result.thread_count:<8} {success_rate:<7.1f}% "
                  f"{result.operations_per_second:<11.1f} {result.avg_response_time*1000:<11.1f}ms "
                  f"{result.memory_usage_mb:<9.1f}MB")
        
        # 性能分析
        print("\n🔍 性能分析:")
        
        # 成功率分析
        stable_results = [r for r in self.results if r.success_count / r.total_operations >= 0.95]
        if stable_results:
            best_stable = max(stable_results, key=lambda x: x.operations_per_second)
            print(f"   ✅ 稳定运行极限: {best_stable.thread_count}线程, {best_stable.operations_per_second:.1f} ops/sec")
        
        # 响应时间分析
        fastest_response = min(self.results, key=lambda x: x.avg_response_time)
        print(f"   ⚡ 最快响应时间: {fastest_response.avg_response_time*1000:.2f}ms ({fastest_response.level})")
        
        # 内存使用分析
        max_memory = max(result.memory_usage_mb for result in self.results)
        print(f"   💾 峰值内存使用: {max_memory:.1f}MB")
        
        # 系统评级
        print("\n🎯 系统评级:")
        if max_ops_per_sec > 1000:
            rating = "🌟🌟🌟🌟🌟 企业级"
        elif max_ops_per_sec > 500:
            rating = "🌟🌟🌟🌟 高性能级"
        elif max_ops_per_sec > 200:
            rating = "🌟🌟🌟 标准级"
        elif max_ops_per_sec > 100:
            rating = "🌟🌟 基础级"
        else:
            rating = "🌟 入门级"
        
        print(f"   {rating}")
        
        # 建议
        print("\n💡 优化建议:")
        if max_concurrent_threads < 20:
            print("   - 考虑优化锁机制，减少线程竞争")
        if any(r.avg_response_time > 0.1 for r in self.results):
            print("   - 考虑增加缓存大小或优化索引")
        if max_memory > 500:
            print("   - 考虑优化内存使用，检查内存泄漏")
        
        print("=" * 60)


def main():
    """主函数"""
    test = MultithreadingStressTest()
    
    try:
        # 初始化
        test.setup()
        
        # 运行所有压力级别测试
        test.run_all_stress_levels()
        
        # 生成最终报告
        test.generate_final_report()
        
        return True
        
    except Exception as e:
        print(f"❌ 测试执行失败: {e}")
        return False
    
    finally:
        test.cleanup()


if __name__ == "__main__":
    try:
        success = main()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n⚠️ 测试被用户中断")
        sys.exit(1)
